import acl
import time
import sys
from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED, ALL_COMPLETED
sys.path.append("..")
from CommonModule.ConstModule import TASK_DONE, TASK_START,\
                                     DEBUG, INFO, WARNING, ERROR,\
                                     ACL_ERROR_NONE
from CommonModule.LogModule import DemoLogger,check_ret

class ThreadPoolInfo():
    def __init__(self):
        self.class_name = self.__class__.__name__
        self.pool = None
        self.total_thread_num = 0
        self.task_thread_num = 0
        self.task_list = []

    def update_thread_pool_info(self, task, task_status):
        if task_status == TASK_START:
            self.task_thread_num += 1
            self.task_list.append(task)
        elif task_status == TASK_DONE:
            self.task_thread_num -= 1
            if self.task_thread_num < 0:
                 DemoLogger.do_log(WARNING, self.class_name, 'update_thread_pool_info', 'task num is low than 0')
            if func in task_list:
                self.task_list.remove(task)
            else:
                DemoLogger.do_log(WARNING, self.class_name, 'update_thread_pool_info', 'the completed task is not in the task_list')
        else:
            pass

    def check_thread_status(self):
        if self.task_thread_num > total_thread_num:
            print("Warning: some tasks are waiting for the available threads in the thread pool: pool: {}, task: {}". \
                  format(total_thread_num, task_thread_num))
        else:
            print("Notice: there are still some available threads in the thread pool: pool: %d, task: %d". \
                  format(total_thread_num, task_thread_num))

class ThreadPool():
    def __init__(self):
        self.class_name = self.__class__.__name__
        self.host_thread_pool_info = ThreadPoolInfo()
        self.device_thread_pool_info = ThreadPoolInfo()
        self.device_func = None

    def host_thread_pool(self, thread_num):
        if thread_num <= 0:
            DemoLogger.do_log(ERROR, self.class_name, 'host_thread_pool', 'thread_num must be larger than 0')
        self.host_thread_pool_info.total_thread_num = thread_num
        self.host_thread_pool_info.pool = ThreadPoolExecutor(max_workers=thread_num)
        if self.host_thread_pool_info.pool is None:
            DemoLogger.do_log(ERROR, self.class_name, 'host_thread_pool', 'fail to create thread pool')
    def combined_function_to_host_thread(self, func, func_args: list):
        task = self.host_thread_pool_info.pool.submit(func, *func_args)
        if task is None:
            DemoLogger.do_log(ERROR, self.class_name, func.__name__, 'fail to submit func to pool')
        self.host_thread_pool_info.update_thread_pool_info(task, TASK_START)

    def add_call_back_func(self, callback_func, task):
        task.add_done_callback(callback_func)

    def host_task_done(self, task):
        info = task.done()
        return info
    
    def wait_host_thread(self, task_list, timeout=None, return_when=ALL_COMPLETED):
        wait(task_list, timeout, return_when)

    def host_task_result(self, func, task):
        result = as_completed(task.result())
        self.host_thread_pool_info.update_thread_pool_info(task, TASK_DONE)
        return result

    def combined_function_to_device_thread(self, func, func_args: list):
        self.device_func = func.__name__
        task, ret = acl.util.start_thread(func, *func_args)
        check_ret(self.class_name, func.__name__, 'fail to star thread by acl.util.start_thread', ret)
        self.device_thread_pool_info.update_thread_pool_info(task, TASK_START)
        return task, ret

    def set_info_when_func_done(self, func, task, thread_pool_info: ThreadPoolInfo):
        if task in thread_pool_info.task_list[func.__name__]:
            self.device_thread_pool_info.update_thread_pool_info(task, TASK_DONE)

    def destroy_host_thread(self):
        if self.host_thread_pool_info.pool is not None:
            self.host_thread_pool_info.pool.shutdown(wait=True)
        self.host_pool_reset()

    def destroy_device_thread(self, thread_id):
        if thread_id is not None:
            ret = acl.util.stop_thread(thread_id)
            check_ret(self.class_name, 'destroy_device_thread','fail to stop thread by acl.util.stop_thread',ret)
            device_pool_reset()

    def host_pool_reset(self):
        self.host_thread_pool_info.pool = None
        self.host_thread_pool_info.total_thread_num = 0
        self.host_thread_pool_info.task_thread_num = 0
        self.host_thread_pool_info.task_list.clear()

    def device_pool_reset(self):
        self.device_thread_pool_info.pool = None
        self.device_thread_pool_info.total_thread_num = 0
        self.device_thread_pool_info.task_thread_num = 0
        self.device_thread_pool_info.task_list.clear()

    #def destroy_all_thread(self):
    #    if self.host_thread_pool_info.pool is not None:
    #        self.host_thread_pool_info.pool.shutdown(wait=TRUE)
    #    for device_thread in self.device_thread_pool_info[self.device_func]:
    #        if device_thread is not None:
    #            ret = self.destroy_device_thread(device_thread)
    #            print("acl stop thread-->ret: %d", ret)
    def  __del__(self):
        self.destroy_host_thread()

